#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Converts json annotations to semantic maps

"""

import os
from operator import attrgetter

import numpy as np
import cv2

from ann_utils import convert_to_path


class ImageAnnotator:
    """

    :Description:

    The Semantic annotation generation classes.

    """

    def __init__(self, url=None, annotation=None, semantic_classes_to_extract=None,
                 semantic_classes_single_plane=None, semantic_classes=None):
        '''

        :param url: The link of the dataset image, part of the anntoation,
        which includes the image name and other details.
        :param annotation: Annotation file name.

        :return: An ImageAnnotator object

        :Description:

        This is the **ImageAnnotator** class constructor.

        '''

        self.url = url

        if '.png' in self.url:
            self.id = self.url.split("/")[-1].rstrip(".png")
        elif '.json' in self.url:
            self.id = self.url.split("/")[-1].rstrip(".json")

        self.annotation = annotation
        self.semantic_classes_to_extract = semantic_classes_to_extract
        self.semantic_classes = semantic_classes
        self.semantic_classes_single_plane = semantic_classes_single_plane

    def save_semantic(self, class_map, dest_file_name=None,
                      paths_dict=None, w=480, h=360):
        '''

        :param class_map: The class map from the json file,
        :param dest_file_name: File name to save the gtLabels.
        :param paths_dict: Output paths.
        :param w: Width of the semantic mask output.
        :param w: Height of the semantic mask output.

        :return: Semantic mask ground truth

        :Description:

        The semantic mask output is generated by parsing the annotation file
        and saved as npz/rgb/gt labels.

        '''

        if not os.path.exists(paths_dict["save_dir_seg"]):
            os.makedirs(paths_dict["save_dir_seg"])
        ann_folder = convert_to_path(paths_dict["save_dir_seg"] +
                                     "/rgbLabels")
        ann_folder_gt = convert_to_path(paths_dict["save_dir_seg"] +
                                        "/gtLabels")
        ann_folder_gt_npz = convert_to_path(paths_dict["save_dir_seg"] +
                                            "/npz")

        if not os.path.exists(ann_folder):
            os.makedirs(ann_folder)
        if not os.path.exists(ann_folder_gt):
            os.makedirs(ann_folder_gt)

        if paths_dict is not None:
            # height, width, channels = img.shape
            height = h
            width = w
            # fill all image in red = background color
            semantic_image = np.full((height, width, 3), (0, 0, 0),
                                     np.uint8)
            # fill all image with 0 = background color
            semantic_image_single_plane = np.full((height, width), 0, np.uint8)

            # Create segmentation mask with hierarchical classes
            # The 'void' class is not needed to create the mask because
            # all classes that are not part of 'road', 'lanemarks' or 'curb'
            # are considered as background
            for obj in self.annotation:
                if class_map[obj.tag] in self.semantic_classes_to_extract:
                    cv2.fillPoly(semantic_image,
                                 [np.array(obj.points).astype(int)],
                                 self.semantic_classes[
                                     class_map[obj.tag]])
                    cv2.fillPoly(semantic_image_single_plane,
                                 [np.array(obj.points).astype(int)],
                                 self.semantic_classes_single_plane[
                                     class_map[obj.tag]])

            semantic_image = cv2.resize(semantic_image, (w, h),
                                        cv2.INTER_NEAREST)
            semantic_image_single_plane = cv2.resize(
                semantic_image_single_plane, (w, h), cv2.INTER_NEAREST)

        if paths_dict is not None:
            cv2.imwrite(convert_to_path(ann_folder + "/" +
                                        self.id + ".png"), semantic_image)
            cv2.imwrite(convert_to_path(ann_folder_gt + "/" + self.id + ".png"),
                        semantic_image_single_plane)


class Instance:
    """

    :Description:

    The Instance level segmentation parser class.

    """

    def __init__(self, json=None):
        '''

        :param json: The annotation json file.

        :return: An Instance object

        :Description:

        This is the **Instance** class constructor. The z-order is appended
        and the relevant tags for objects are appended

        '''
        if json is not None:
            self.z_order = json["z_order"]
            self.tag = json["tags"][0]
            if len(json["tags"]) != 1:
                print("unexpected amount of TAGS: ", json["tags"])
            self.points = json["segmentation"]
            self.id = json["id"]


def parse_annotation(ann_list):
    '''

    :param ann_list: The list of annotations in a particular annotation file.

    :return: Annotation list, sorted based on z-order

    :Description:

    The annotations are sorted based on z-order and returned.

    '''
    ann = []
    for elem in ann_list:
        if len(elem["tags"]) == 1:
            ann.append(Instance(json=elem))
    ann = sorted(ann, key=attrgetter('z_order'))
    return ann
